"""

This script will take a list of user year clustered user vectors, the output of classifyVectors.py

And output statistics about

for each cohor, and for each lifespan group of that cohort, inlcuding 
amount of work, by cluster, by dat
proportion of work, by cluster, by date

""" 

import os
import sys
import datetime

import wp_datetime as dt
import wp_namespace as ns

def main():

        inputDir = sys.argv[1]
        outputDir = sys.argv[2]
	clusters = int(sys.argv[3])

	START_COHORT = 2003
	END_COHORT = 2012

	END_YEAR = 2012
	CURRENT_CUTOFF = 2011

	work = {}
	current_work = {}
	
	for year in range(START_COHORT, END_COHORT + 1):

        	work[year] = {}
		current_work[year] = {}

                for offset in range (0, END_YEAR - year + 1):

                        work[year][offset] = {}
	


        inputFileList = os.listdir(inputDir)

        for inputFileName in inputFileList:

                inputFilePath = inputDir+'/'+inputFileName

                if os.path.isfile(inputFilePath):
                        processFile(inputFilePath, work, current_work, CURRENT_CUTOFF, clusters)

	current_work_prop = {}
	for cohort in current_work:
		current_work_prop[cohort] = findProportion(current_work[cohort], clusters)

	work_prop = {}
	for cohort in work:
		work_prop[cohort] = {}
		for lifespan in work[cohort]:
			work_prop[cohort][lifespan] = findProportion(work[cohort][lifespan], clusters)

	
	for cohort in current_work:
		
			printStats(current_work[cohort], outputDir + '/' + 'cohort'+str(cohort)+'current_count.stat', clusters)
			printStats(current_work_prop[cohort], outputDir + '/' + 'cohort'+str(cohort)+'current_prop.stat', clusters)
	
	for cohort in work:
		for lifespan in work[cohort]:

			# Make sure its not the last year, because it is always current
			if END_YEAR - lifespan != cohort:
				printStats(work[cohort][lifespan], outputDir + '/' + 'cohort'+str(cohort)+'lifespan'+str(lifespan)+'_count.stat', clusters)
                        	printStats(work_prop[cohort][lifespan], outputDir + '/' + 'cohort'+str(cohort)+'lifespan'+str(lifespan)+'_propt.stat', clusters)



def printStats(work, path, clusters):

	header = ['Date']

	header.extend([str(num) for num in range(1, clusters + 1)])

	outfile = open(path, 'w')

	outfile.write('\t'.join(header) + '\n')

	for date in sorted(work):

		out = [str(date)]

		for cluster in range(1, clusters + 1):

			out.append(str(work[date][cluster]))

		outfile.write('\t'.join(out) + '\n')

	outfile.close()

def findProportion(data, clusters):

	prop = {}
	for date in data:
		initWorkTable(prop, date, clusters)

		total_work = 0
		for cluster in data[date]:
			total_work += data[date][cluster]

		for cluster in data[date]:
			if total_work == 0:
				prop[date][cluster] = 0
			else:
				prop[date][cluster] = float(data[date][cluster]) / float(total_work)

	return prop
			

def processFile(inputFilePath, work, current_work, CURRENT_CUTOFF, clusters):

        print(timestamp() + " Processing " + inputFilePath)

        inputFileHandle = open(inputFilePath, 'r')

	userdata = []

        for line in inputFileHandle:

                fields = parseLine(line)

		if userdata and fields['username'] != userdata[0]['username']:
			addUser(work, current_work, userdata, CURRENT_CUTOFF, clusters)
			userdata = []

		userdata.append(fields)

	addUser(work, current_work, userdata, CURRENT_CUTOFF, clusters)

	inputFileHandle.close()

        print(timestamp() + " Finished processing " + inputFilePath)


def addUser(work, current_work, userdata, cutoff, clusters):

	if userdata[-1]['year'] < cutoff:
		for line in userdata:
			addWork(work, line, clusters)
	else:
		for line in userdata:
			addWorkToCurrent(current_work, line, clusters)

def addWorkToCurrent(worktable, fields, clusters):

	cohort = fields['cohort']
	cluster = fields['cluster']

	#date = simpleDate(fields['timestamp'])
	date = str(int(fields['cohort']) + int(fields['offset']))       # only use year...
	if not date in worktable[cohort]:
		initWorkTable(worktable[cohort], date, clusters)

        worktable[cohort][date][cluster] += 1


def addWork(worktable, fields, clusters):

        cohort = fields['cohort']
        lifespan = fields['lifespan']
	cluster = fields['cluster']	

	#date = simpleDate(fields['timestamp'])   
	date = str(int(fields['cohort']) + int(fields['offset']))       # only use year...
	if not date in worktable[cohort][lifespan]:
  		initWorkTable(worktable[cohort][lifespan], date, clusters)

       	worktable[cohort][lifespan][date][cluster] += 1

def parseLine(line):

        tokens = line.strip().split('\t')

        fields = {}
        fields['username'] = tokens[0]
        fields['cohort'] = int(tokens[1])
        fields['lifespan'] = int(tokens[2])
        fields['offset'] = int(tokens[3])
        fields['year'] = int(tokens[4])
        fields['cluster'] = int(tokens[5])

        return fields

def initWorkTable(work, date, clusters):
        """
        Inits a count for a date in the work table. Must be the simple date returned from simpleDate(), not a timestamp.
        """

        work[date] = {}
        work[date] = {num: 0.0 for num in range(1, clusters + 1)}


def simpleDate(timestamp):

        return timestamp.split('T')[0]

def timestamp():

        return datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')


if __name__ == '__main__':
        main()

